from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.wait import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver import ActionChains

import base64
import random
import time
from datetime import datetime

import numpy as np

import skimage
from skimage import io
from skimage.color import rgba2rgb,rgb2gray
from skimage.util import compare_images
from skimage.measure import find_contours
from skimage.filters import threshold_otsu

import matplotlib.pyplot as plt

# 取消显示滑块
# js_slice = "document.getElementsByClassName('geetest_canvas_slice')[0].style.display='none'" 
# 取消显示拼图
# js_jigsaw = "document.getElementsByClassName('geetest_canvas_fullbg')[0].style.display='block'"

# 浏览器
driver = webdriver.Chrome()

# geetest测试页面
# target_url = "https://www.geetest.com/Register"
target_url = "https://www.geetest.com/demo/slide-bind.html"

# 人工鼠标轨迹记录
with open("mouse/mouse_track.txt") as f:
    track_data = f.read()

track_data = track_data.split("\n\n")

# 轨迹字典
# 键：拖动总长度
# 值：轨迹数组，每个数字是鼠标横向移动一次的像素数
track_map = {}

for line in track_data:
    line_list = [int(x) for x in line.split(",")]
    k = line_list[-1]
    track_map[k] = np.array(line_list[1:]) - np.array(line_list[:-1])

class geetestHacker:
    
    def __init__(self, driver, target_url):
        self.driver = driver
        self.target_url = target_url
        
        # 打开验证码页面
        self.driver.get(self.target_url)
        
        self.wait = WebDriverWait(self.driver, 3)
        
        self.jigsaw_imgname = None
        self.fullbg_imgname = None
        self.x_offset = 0
        self.track = []
        
    def open_captcha(self):
        
        time.sleep(1)
        # 等待按钮可点击
        
        # 测试页面
        submit_btn = self.wait.until(EC.element_to_be_clickable((By.ID, 'btn')))
        # 官网注册页面
        # submit_btn = self.wait.until(EC.element_to_be_clickable((By.CLASS_NAME, 'sendCode')))
        
        # 点击“验证按钮”
        submit_btn.click()
        # 等待滑块出现
        canvas = self.wait.until(EC.presence_of_element_located((By.CLASS_NAME, 'geetest_slicebg')))
        
    def get_image(self):
        
        # 获取图片地址
        js_getimg_jigsaw = "return document.getElementsByClassName('geetest_canvas_bg')[0].toDataURL('image/png')"
        js_getimg_fullbg = "return document.getElementsByClassName('geetest_canvas_fullbg')[0].toDataURL('image/png')"
        
        # 生成时间+随机数文件名后缀
        suffix = datetime.now().strftime("%Y%m%d%H%M%S") + \
                 "{:>03}".format(random.randint(0,100))
        
        # 获取带有锯齿阴影的图片
        jigsaw_imgpath = driver.execute_script(js_getimg_jigsaw)
        jigsaw_imgdata = base64.b64decode(jigsaw_imgpath.split(",")[1])

        self.jigsaw_imgname = "images/jigsaw_" + \
                          suffix + \
                         ".png"

        with open(self.jigsaw_imgname, "wb") as f:
            f.write(jigsaw_imgdata)
            
        # 获取完整图片
        fullbg_imgpath = driver.execute_script(js_getimg_fullbg)
        fullbg_imgdata = base64.b64decode(fullbg_imgpath.split(",")[1])

        self.fullbg_imgname = "images/fullbg_" + \
                               suffix + \
                              ".png"

        with open(self.fullbg_imgname, "wb") as f:
            f.write(fullbg_imgdata)
        
        
    def compute_offset(self):
        # 读取图片
        img_jigsaw = io.imread(self.jigsaw_imgname)
        img_fullbg = io.imread(self.fullbg_imgname)
        # 灰度化
        img_jigsaw = rgb2gray(rgba2rgb(img_jigsaw))
        img_fullbg = rgb2gray(rgba2rgb(img_fullbg))
        # 求差值
        img_diff = compare_images(img_jigsaw, img_fullbg)
        
        # 二值化
        binary = img_diff > threshold_otsu(img_diff)
        # 最大contour
        contours = find_contours(binary, 0.5)
        # 最左侧横坐标
        self.x_offset = contours[0][:,1].min() - 6
        
    
    def gen_track(self):
        
        # 查找长度最接近的轨迹
        def find_closest():
            return sorted(list(track_map.keys()), key=lambda x: abs(self.x_offset-x))[0]
        
        self.track = track_map[find_closest()]        
        
    def slide(self):
        # 获取滑块对象
        slider = self.wait.until(EC.element_to_be_clickable((By.CLASS_NAME, 'geetest_slider_button')))

        # 点击并按住
        ActionChains(self.driver).click_and_hold(slider).perform()
        
        # 按照轨迹数组运动
        for x in self.track:
            ActionChains(self.driver).move_by_offset(xoffset=x, yoffset=0).perform()

        time.sleep(0.5)
        # 松开滑块
        ActionChains(self.driver).release().perform()

hacker = geetestHacker(driver, target_url)

hacker.open_captcha()

hacker.get_image()
hacker.compute_offset()
hacker.gen_track()
hacker.slide()

driver.close()